import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import os
from collections import defaultdict
from cec2017.functions import all_functions
from matplotlib.colors import LinearSegmentedColormap
import matplotlib as mpl
from cec2017.basic import ackley as ackley
func_id = 9

# 定义测试函数 (Sphere函数)
def sphere_function(x):
    # # 确保x是正确的形状以适合all_functions[10]
    if len(x.shape) == 1:
        # 如果是单个向量，确保它是2D形状 (1,dim)
        x = (x.reshape(1, -1))  # 返回标量
    return  ackley(x)

# 定义DE算法类
class DifferentialEvolution:
    def __init__(self, fitness_func, dim=10, pop_size=50, F=0.5, CR=0.7, 
                 bounds=(-100, 100), max_iter=100, p_value=0.5, initial_pop=None):
        self.fitness_func = fitness_func
        self.dim = dim
        self.pop_size = pop_size
        self.F = F
        self.CR = CR
        self.bounds = bounds
        self.max_iter = max_iter
        self.p_value = p_value
        
        # 初始化种群 - 如果提供了初始种群则使用它，否则随机生成
        if initial_pop is not None:
            self.population = initial_pop.copy()
        else:
            self.population = np.random.uniform(
                low=self.bounds[0], high=self.bounds[1], 
                size=(self.pop_size, self.dim)
            )
        
        # 评估初始种群
        self.fitness = self.fitness_func(self.population)
        
        # 记录全局最优
        self.best_idx = np.argmin(self.fitness)
        self.best_solution = self.population[self.best_idx].copy()
        self.best_fitness = self.fitness[self.best_idx]
        
        # 记录历史数据
        self.fitness_history = [self.best_fitness]
        self.center_distances = []
        
        # 存储所有评估过的位置 - 每代都添加新评估的点
        self.all_positions = self.population.copy()
        self.all_fitnesses = self.fitness.copy()
        
        # 计算初始种群中心距离
        self._update_center_distance()
    
    def _update_center_distance(self):
        """计算并记录种群到中心点的平均曼哈顿距离"""
        center = np.mean(self.population, axis=0)
        manhattan_distances = np.mean(np.abs(self.population - center), axis=1)
        avg_distance = np.mean(manhattan_distances)
        self.center_distances.append(avg_distance)
    
    def _select_parents(self, i):
        """为第i个个体选择父代索引"""
        idxs = list(range(self.pop_size))
        idxs.remove(i)
        return np.random.choice(idxs, 3, replace=False)
    
    def run(self):
        """运行DE算法"""
        for gen in range(self.max_iter):
            # 对种群按适应度排序（最小化问题）
            sorted_indices = np.argsort(self.fitness)
            self.population = self.population[sorted_indices]
            self.fitness = self.fitness[sorted_indices]
            
            # 确定使用开发算子的数量
            exploit_count = int(self.pop_size * self.p_value)
            
            # 创建新一代种群
            new_population = np.zeros_like(self.population)
            new_fitness = np.zeros_like(self.fitness)
            
            # 存储这一代评估的所有试验向量
            gen_trials = []
            gen_fitness = []
            
            for i in range(self.pop_size):
                # 选择父代
                r1, r2, r3 = self._select_parents(i)
                
                # 创建试验向量 (区分勘探和开发算子)
                if i < exploit_count:  # 开发算子
                    mutant = self.population[i] + self.F * (self.best_solution - self.population[i]) + self.F * (self.population[r1] - self.population[r2])
                else:  # 勘探算子
                    mutant = self.population[r1] + self.F * (self.population[r2] - self.population[r3])
                
                # 边界处理
                mutant = np.clip(mutant, self.bounds[0], self.bounds[1])
                
                # 交叉
                cross_points = np.random.rand(self.dim) <= self.CR
                if not np.any(cross_points):
                    cross_points[np.random.randint(0, self.dim)] = True
                
                trial = np.where(cross_points, mutant, self.population[i])
                
                # 评估新解
                trial_fitness = self.fitness_func(trial)
                
                # 存储评估的位置和适应度
                gen_trials.append(trial)
                gen_fitness.append(trial_fitness)
                
                # 选择
                if trial_fitness <= self.fitness[i]:
                    new_population[i] = trial
                    new_fitness[i] = trial_fitness
                else:
                    new_population[i] = self.population[i]
                    new_fitness[i] = self.fitness[i]
            
            # 更新种群
            self.population = new_population
            self.fitness = new_fitness
            
            # 更新全局最优
            best_idx = np.argmin(self.fitness)
            if self.fitness[best_idx] < self.best_fitness:
                self.best_fitness = self.fitness[best_idx]
                self.best_solution = self.population[best_idx].copy()
            
            # 记录历史数据
            self.fitness_history.append(self.best_fitness)
            self._update_center_distance()
            
            # 添加本代评估的所有点
            self.all_positions = np.vstack((self.all_positions, np.array(gen_trials)))
            self.all_fitnesses = np.append(self.all_fitnesses, gen_fitness)
    
    def get_plot_data(self):
        """返回绘图数据：前两维的位置和适应度"""
        return self.all_positions[:, :2], self.all_fitnesses

# 辅助函数 - 计算所有个体间的平均距离
def calculate_avg_pairwise_distance(positions, max_samples=5000):
    """计算所有点之间的平均欧几里德距离（如果点太多则采样）"""
    n = positions.shape[0]
    
    # 如果点太多，进行采样
    if n > max_samples:
        indices = np.random.choice(n, max_samples, replace=False)
        sampled_positions = positions[indices]
    else:
        sampled_positions = positions
    
    # 使用scipy的pdist高效计算所有点对之间的距离
    from scipy.spatial.distance import pdist
    distances = pdist(sampled_positions)  # 计算压缩的距离矩阵
    avg_distance = np.mean(distances)
    
    return avg_distance

# 设置实验参数
dim = 2  # 绘图只显示前两维，所以使用2维
pop_size = 200
bounds = (-100, 100)
max_iter = 50
p_values = [0.1, 0.45, 0.9]  # 两种p值

# 使用新的随机种子以获得不同的结果
np.random.seed(87765)

# 创建结果目录
os.makedirs("results", exist_ok=True)

# 进行实验
results = {}
for p in p_values:
    print(f"Running experiment with p = {p}")
    de = DifferentialEvolution(
        fitness_func=sphere_function,
        dim=dim,
        pop_size=pop_size,
        bounds=bounds,
        max_iter=max_iter,
        p_value=p
    )
    de.run()
    positions, fitnesses = de.get_plot_data()
    results[p] = {
        'positions': positions,
        'fitnesses': fitnesses,
        'fitness_history': de.fitness_history,
        'center_distances': de.center_distances
    }
    print(f"Experiment with p = {p} completed.")

# 设置绘图风格 - 学术风格
plt.rcParams.update({
    'font.family': 'serif',
    'font.size': 12,
    'axes.linewidth': 1.5,
    'axes.grid': True,
    'grid.alpha': 0.3,
    'grid.linestyle': '--',
    'figure.dpi': 300,
    'savefig.dpi': 300,
    'savefig.format': 'pdf',
    'savefig.bbox': 'tight'
})

# 设置不同的蓝色梯度
blue_colors = {
    p_values[0]: '#2E4053',  # 浅蓝色 - 强勘探策略
    p_values[1]: '#B8860B',  # 中蓝色 - 中等勘探策略
    p_values[2]: '#228B22'   # 深蓝色 - 强开发策略
}

# 创建散点图 - 展示搜索模式
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
fig.suptitle(
    'Influence of Control Parameter \'p\' on DE Population Distribution \n During 2D Ackley Function Optimization', 
    fontsize=16, fontweight='bold', y=1.02
)

for i, p in enumerate(p_values):
    positions = results[p]['positions']
    fitnesses = results[p]['fitnesses']
    
    # 计算平均成对距离
    avg_distance = calculate_avg_pairwise_distance(positions)
    
    # 设置正方形形状
    axes[i].set_aspect('equal')
    
    # 使用对应的蓝色进行可视化
    scatter = axes[i].scatter(
        positions[:, 0], 
        positions[:, 1],
        color=blue_colors[p],
        alpha=0.5,
        s=10,
        edgecolor='none'
    )
    
    # 设置标题和标签，包含平均距离信息
    strategy = 'Exploration Strategy' if p < 0.5 else 'Exploitation Strategy'
    title = f'{strategy} (p={p})\nAvg. Pairwise Distance: {avg_distance:.2f}'
    axes[i].set_title(title, fontsize=14, fontweight='bold', pad=10)
    axes[i].set_xlabel('Dimension 1', fontweight='bold')
    axes[i].set_ylabel('Dimension 2', fontweight='bold')
    
    # 设置坐标轴范围
    axes[i].set_xlim(bounds)
    axes[i].set_ylim(bounds)
    
    # 添加全局最优点的标记
    best_idx = np.argmin(fitnesses)
    axes[i].scatter(
        positions[best_idx, 0],
        positions[best_idx, 1],
        marker='*',
        s=200,
        color='red',
        edgecolor='black',
        linewidth=1.5,
        label='Global Best'
    )
    
    axes[i].legend(loc='upper right', frameon=True)

plt.tight_layout() 
if not os.path.exists('graph/visual_exploration'):
    os.makedirs('graph/visual_exploration')
plt.savefig('graph/visual_exploration/de_search_patterns_scatter.png', dpi=300, bbox_inches='tight')

plt.savefig('graph/visual_exploration/de_search_patterns_scatter.eps', dpi=300, bbox_inches='tight')
  